package com.king;

import java.io.IOException;
import java.lang.instrument.Instrumentation;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;

public class NacosAgent {

    /** 被转换的类 */
//    public static final String TRANSFORM_CLASS = "org.example.agent.AgentTest";

//    /** 静态加载。Java agent指定的premain方法，会在main方法之前被调用 */
//    public static void premain(String args, Instrumentation instrumentation) {
//        System.out.println("premain start!");
//        addTransformer(instrumentation);
//        System.out.println("premain end!");
//    }
//
//    /** 动态加载。Java agent指定的premain方法，会在main方法之前被调用 */
//    public static void agentmain(String args, Instrumentation instrumentation) {
//        System.out.println("agentmain start!");
//        addTransformer(instrumentation);
//        Class<?>[] classes = instrumentation.getAllLoadedClasses();
//        if (classes != null){
//            for (Class<?> c: classes) {
//                if (c.isInterface() ||c.isAnnotation() ||c.isArray() ||c.isEnum()){
//                    continue;
//                }
//                if (c.getName().equals(TRANSFORM_CLASS)) {
//                    try {
//                        System.out.println("retransformClasses start, class: " + c.getName());
//                        /*
//                         * retransformClasses()对JVM已经加载的类重新触发类加载。使用的就是上面注册的Transformer。
//                         * retransformClasses()可以修改方法体，但是不能变更方法签名、增加和删除方法/类的成员属性
//                         */
//                        instrumentation.retransformClasses(c);
//                        System.out.println("retransformClasses end, class: " + c.getName());
//                    } catch (UnmodifiableClassException e) {
//                        System.out.println("retransformClasses error, class: " + c.getName() + ", ex:" + e);
//                        e.printStackTrace();
//                    }
//                }
//            }
//        }
//        System.out.println("agentmain end!");
//    }
//
//    private static void addTransformer (Instrumentation instrumentation) {
//        /* Instrumentation提供的addTransformer方法，在类加载时会回调ClassFileTransformer接口 */
//        instrumentation.addTransformer(new ClassFileTransformer() {
//            public byte[] transform(ClassLoader l, String className, Class<?> c, ProtectionDomain pd, byte[] b){
//                try {
//                    className = className.replace("/", ".");
//                    if (className.equals(TRANSFORM_CLASS)) {
//                        final ClassPool classPool = ClassPool.getDefault();
//                        final CtClass clazz = classPool.get(TRANSFORM_CLASS);
//
//                        for (CtMethod method : clazz.getMethods()) {
//                            /*
//                             * Modifier.isNative(methods[i].getModifiers())过滤本地方法,否则会报
//                             * javassist.CannotCompileException: no method body  at javassist.CtBehavior.addLocalVariable()
//                             * 报错原因如下
//                             * 来自Stack Overflow网友解答
//                             * Native methods cannot be instrumented because they have no bytecodes.
//                             * However if native method prefix is supported ( Transformer.isNativeMethodPrefixSupported() )
//                             * then you can use Transformer.setNativeMethodPrefix() to wrap a native method call inside a non-native call
//                             * which can then be instrumented
//                             */
//                            if (Modifier.isNative(method.getModifiers())) {
//                                continue;
//                            }
//
//                            method.insertBefore("System.out.println(\"" + clazz.getSimpleName() + "."
//                                    + method.getName() + " start.\");");
//                            method.insertAfter("System.out.println(\"" + clazz.getSimpleName() + "."
//                                    + method.getName() + " end.\");", false);
//                        }
//
//                        return clazz.toBytecode();
//                    }
//                } catch (Exception e) {
//                    e.printStackTrace();
//                }
//
//                return null;
//            }
//        }, true);
//    }

    private static List<String> configLines = new ArrayList<>();

    public static void premain(String agentArgs, Instrumentation inst) {
        // 读取配置文件
        if (agentArgs != null && !agentArgs.isEmpty()) {
            loadConfigFromFile(agentArgs);
        } else {
            // 尝试从当前jar包目录读取配置文件
            String jarFilePath = NacosAgent.class.getProtectionDomain().getCodeSource().getLocation().getPath();
            String jarDir = jarFilePath.substring(0, jarFilePath.lastIndexOf('/') + 1);
            String configPath = jarDir + "nacosConfig.txt";
            if (configPath.startsWith("/")){
                configPath = configPath.replaceFirst("/", "");
            }
            loadConfigFromFile(configPath);
        }

        String classNameTarget = "com.alibaba.nacos.client.naming.remote.NamingClientProxyDelegate";
        inst.addTransformer((loader, className, classBeingRedefined, protectionDomain, classfileBuffer) -> {
            if (classNameTarget.equals(className)) {
                // 使用Javassist库对类进行转换
                javassist.ClassPool pool = javassist.ClassPool.getDefault();
                try {
                    javassist.CtClass cc = pool.get(classNameTarget);
                    javassist.CtMethod m = cc.getDeclaredMethod("queryInstancesOfService");
                    StringBuilder codeBlock = new StringBuilder("   Object result = proceed($$); " +
                                    // 解析配置文件
                                    "   Map<String, String> configMap = new java.util.HashMap<>();");
                    for (int i = 1; i < configLines.size(); i++) {
                        String[] splitStr = configLines.get(i).split(" ");
                        if (splitStr.length == 3){
                            codeBlock.append("   configMap.put(\"").append(splitStr[0]).append("\", \"").append(splitStr[1]).append(",").append(splitStr[2]).append("\");");
                        }
                    }
                    codeBlock.append("   if(result != null && result instanceof com.alibaba.nacos.api.naming.pojo.ServiceInfo) {" +
                            "       com.alibaba.nacos.api.naming.pojo.ServiceInfo instance = (com.alibaba.nacos.api.naming.pojo.ServiceInfo)result;" +
                            "         if (configMap.containsKey(instance.getName())){\n" +
                            "            String configStr = configMap.get(instance.getName());\n" +
                            "            String[] split = configStr.split(\",\");\n" +
                            "            List<com.alibaba.nacos.api.naming.pojo.Instance> hosts = instance.getHosts();\n" +
                            "            hosts.forEach(item ->{\n" +
                            "                item.setIp(split[0]);\n" +
                            "                item.setPort(Integer.parseInt(split[1]));\n" +
                            "            }\n" +
                            "        });" +
                            "   }" +
                            "   return result;");
                    String block = codeBlock.toString();
                    System.out.println("代理代码块：" + block);
                    m.insertBefore(block);
                    return cc.toBytecode();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
            return classfileBuffer;
        });
    }

    private static void loadConfigFromFile(String filePath) {
        try {
            List<String> lines = Files.readAllLines(Paths.get(filePath));
            configLines.addAll(lines);
        } catch (IOException e) {
            System.err.println("Failed to read config file: " + filePath);
            e.printStackTrace();
        }
    }

    public static List<String> getConfigLines() {
        return configLines;
    }
}
